'''
https://leetcode.cn/problems/paths-in-matrix-whose-sum-is-divisible-by-k/description/
'''
from functools import cache
from typing import List


class Solution:
    def numberOfPaths(self, grid: List[List[int]], k: int) -> int:
        m, n = len(grid), len(grid[0])
        MOD = 10 ** 9 + 7

        # i,j 当前位置，v: 在i,j位置希望的余数
        @cache
        def f(i, j, remainder):
            if i == m - 1 and j == n - 1:
                return 1 if grid[i][j] % k == remainder else 0
            # want_other_remainder - grid[i][j] % k
            want_other_remainder = (k + remainder - grid[i][j] % k) % k
            res = 0
            if i + 1 < m:
                res += f(i + 1, j, want_other_remainder)
            if j + 1 < n:
                res += f(i, j + 1, want_other_remainder)
            return res % MOD

        return f(0, 0, 0)

    # dp 打表
    def numberOfPaths2(self, grid: List[List[int]], k: int) -> int:
        m, n = len(grid), len(grid[0])
        MOD = 10 ** 9 + 7
        dp = [[[0] * k for _ in range(n)] for _ in range(m)]
        dp[-1][-1][grid[-1][-1] % k] = 1
        for i in range(m - 2, -1, -1):
            for remainder in range(k):
                want_other_remainder = (k + remainder - grid[i][-1] % k) % k
                dp[i][-1][remainder] = dp[i + 1][-1][want_other_remainder]
        for j in range(n - 2, -1, -1):
            for remainder in range(k):
                want_other_remainder = (k + remainder - grid[-1][j] % k) % k
                dp[-1][j][remainder] = dp[-1][j + 1][want_other_remainder]
        # 第一维度依赖后边的，第二维度依赖后边的，第三维度随便依赖的单肯定不是当前层了
        for i in range(m - 2, -1, -1):
            for j in range(n - 2, -1, -1):
                for remainder in range(k):
                    want_other_remainder = (k + remainder - grid[i][j] % k) % k
                    res = 0
                    if i + 1 < m:
                        res += dp[i + 1][j][want_other_remainder]
                    if j + 1 < n:
                        res += dp[i][j + 1][want_other_remainder]
                    dp[i][j][remainder] = res % MOD
        return dp[0][0][0]


# k = 10
# for cur in range(20):
#     for remainder in range(k):
#         # print(f'当前数: {cur}, 期望当前的余数: {remainder}, 期望后边的余数: {remainder - cur % k}')
#         print(f'当前数: {cur}, 期望当前的余数: {remainder}, 期望后边的余数: {(k + remainder - cur % k) % k}')

grid = [[5, 2, 4], [3, 0, 5], [0, 7, 2]]
k = 3
print(Solution().numberOfPaths2(grid, k))
